[3.6] bpo-30695: Add set_nomemory(start, stop) to _testcapi (GH-2406)… · python/cpython@aaf6a3d (original) (raw)
`@@ -3437,6 +3437,130 @@ test_pyobject_setallocators(PyObject *self)
`
3437
3437
`return test_setallocators(PYMEM_DOMAIN_OBJ);
`
3438
3438
`}
`
3439
3439
``
``
3440
`+
/* Most part of the following code is inherited from the pyfailmalloc project
`
``
3441
`+
- written by Victor Stinner. */
`
``
3442
`+
static struct {
`
``
3443
`+
int installed;
`
``
3444
`+
PyMemAllocatorEx raw;
`
``
3445
`+
PyMemAllocatorEx mem;
`
``
3446
`+
PyMemAllocatorEx obj;
`
``
3447
`+
} FmHook;
`
``
3448
+
``
3449
`+
static struct {
`
``
3450
`+
int start;
`
``
3451
`+
int stop;
`
``
3452
`+
Py_ssize_t count;
`
``
3453
`+
} FmData;
`
``
3454
+
``
3455
`+
static int
`
``
3456
`+
fm_nomemory(void)
`
``
3457
`+
{
`
``
3458
`+
FmData.count++;
`
``
3459
`+
if (FmData.count > FmData.start &&
`
``
3460
`+
(FmData.stop <= 0 || FmData.count <= FmData.stop)) {
`
``
3461
`+
return 1;
`
``
3462
`+
}
`
``
3463
`+
return 0;
`
``
3464
`+
}
`
``
3465
+
``
3466
`+
static void *
`
``
3467
`+
hook_fmalloc(void *ctx, size_t size)
`
``
3468
`+
{
`
``
3469
`+
PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;
`
``
3470
`+
if (fm_nomemory()) {
`
``
3471
`+
return NULL;
`
``
3472
`+
}
`
``
3473
`+
return alloc->malloc(alloc->ctx, size);
`
``
3474
`+
}
`
``
3475
+
``
3476
`+
static void *
`
``
3477
`+
hook_fcalloc(void *ctx, size_t nelem, size_t elsize)
`
``
3478
`+
{
`
``
3479
`+
PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;
`
``
3480
`+
if (fm_nomemory()) {
`
``
3481
`+
return NULL;
`
``
3482
`+
}
`
``
3483
`+
return alloc->calloc(alloc->ctx, nelem, elsize);
`
``
3484
`+
}
`
``
3485
+
``
3486
`+
static void *
`
``
3487
`+
hook_frealloc(void *ctx, void *ptr, size_t new_size)
`
``
3488
`+
{
`
``
3489
`+
PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;
`
``
3490
`+
if (fm_nomemory()) {
`
``
3491
`+
return NULL;
`
``
3492
`+
}
`
``
3493
`+
return alloc->realloc(alloc->ctx, ptr, new_size);
`
``
3494
`+
}
`
``
3495
+
``
3496
`+
static void
`
``
3497
`+
hook_ffree(void *ctx, void *ptr)
`
``
3498
`+
{
`
``
3499
`+
PyMemAllocatorEx *alloc = (PyMemAllocatorEx *)ctx;
`
``
3500
`+
alloc->free(alloc->ctx, ptr);
`
``
3501
`+
}
`
``
3502
+
``
3503
`+
static void
`
``
3504
`+
fm_setup_hooks(void)
`
``
3505
`+
{
`
``
3506
`+
PyMemAllocatorEx alloc;
`
``
3507
+
``
3508
`+
if (FmHook.installed) {
`
``
3509
`+
return;
`
``
3510
`+
}
`
``
3511
`+
FmHook.installed = 1;
`
``
3512
+
``
3513
`+
alloc.malloc = hook_fmalloc;
`
``
3514
`+
alloc.calloc = hook_fcalloc;
`
``
3515
`+
alloc.realloc = hook_frealloc;
`
``
3516
`+
alloc.free = hook_ffree;
`
``
3517
`+
PyMem_GetAllocator(PYMEM_DOMAIN_RAW, &FmHook.raw);
`
``
3518
`+
PyMem_GetAllocator(PYMEM_DOMAIN_MEM, &FmHook.mem);
`
``
3519
`+
PyMem_GetAllocator(PYMEM_DOMAIN_OBJ, &FmHook.obj);
`
``
3520
+
``
3521
`+
alloc.ctx = &FmHook.raw;
`
``
3522
`+
PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &alloc);
`
``
3523
+
``
3524
`+
alloc.ctx = &FmHook.mem;
`
``
3525
`+
PyMem_SetAllocator(PYMEM_DOMAIN_MEM, &alloc);
`
``
3526
+
``
3527
`+
alloc.ctx = &FmHook.obj;
`
``
3528
`+
PyMem_SetAllocator(PYMEM_DOMAIN_OBJ, &alloc);
`
``
3529
`+
}
`
``
3530
+
``
3531
`+
static void
`
``
3532
`+
fm_remove_hooks(void)
`
``
3533
`+
{
`
``
3534
`+
if (FmHook.installed) {
`
``
3535
`+
FmHook.installed = 0;
`
``
3536
`+
PyMem_SetAllocator(PYMEM_DOMAIN_RAW, &FmHook.raw);
`
``
3537
`+
PyMem_SetAllocator(PYMEM_DOMAIN_MEM, &FmHook.mem);
`
``
3538
`+
PyMem_SetAllocator(PYMEM_DOMAIN_OBJ, &FmHook.obj);
`
``
3539
`+
}
`
``
3540
`+
}
`
``
3541
+
``
3542
`+
static PyObject*
`
``
3543
`+
set_nomemory(PyObject *self, PyObject *args)
`
``
3544
`+
{
`
``
3545
`+
/* Memory allocation fails after 'start' allocation requests, and until
`
``
3546
`+
- 'stop' allocation requests except when 'stop' is negative or equal
`
``
3547
`+
- to 0 (default) in which case allocation failures never stop. */
`
``
3548
`+
FmData.count = 0;
`
``
3549
`+
FmData.stop = 0;
`
``
3550
`+
if (!PyArg_ParseTuple(args, "i|i", &FmData.start, &FmData.stop)) {
`
``
3551
`+
return NULL;
`
``
3552
`+
}
`
``
3553
`+
fm_setup_hooks();
`
``
3554
`+
Py_RETURN_NONE;
`
``
3555
`+
}
`
``
3556
+
``
3557
`+
static PyObject*
`
``
3558
`+
remove_mem_hooks(PyObject *self)
`
``
3559
`+
{
`
``
3560
`+
fm_remove_hooks();
`
``
3561
`+
Py_RETURN_NONE;
`
``
3562
`+
}
`
``
3563
+
3440
3564
`PyDoc_STRVAR(docstring_empty,
`
3441
3565
`""
`
3442
3566
`);
`
`@@ -4318,6 +4442,10 @@ static PyMethodDef TestMethods[] = {
`
4318
4442
` (PyCFunction)test_pymem_setallocators, METH_NOARGS},
`
4319
4443
` {"test_pyobject_setallocators",
`
4320
4444
` (PyCFunction)test_pyobject_setallocators, METH_NOARGS},
`
``
4445
`+
{"set_nomemory", (PyCFunction)set_nomemory, METH_VARARGS,
`
``
4446
`+
PyDoc_STR("set_nomemory(start:int, stop:int = 0)")},
`
``
4447
`+
{"remove_mem_hooks", (PyCFunction)remove_mem_hooks, METH_NOARGS,
`
``
4448
`+
PyDoc_STR("Remove memory hooks.")},
`
4321
4449
` {"no_docstring",
`
4322
4450
` (PyCFunction)test_with_docstring, METH_NOARGS},
`
4323
4451
` {"docstring_empty",
`