[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

`+

`

``

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

`+

`

``

3547

`+

`

``

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",

`