Skip to content

Commit da64667

Browse files
ngoldbaumngoldbaum
authored
Merge pull request #94 from ngoldbaum/alloc-mutex
Use a mutex to lock heap memory access
2 parents e663f0f + 44d9943 commit da64667

File tree

7 files changed

+707
-282
lines changed

7 files changed

+707
-282
lines changed

‎stringdtype/stringdtype/src/casts.c

Lines changed: 121 additions & 27 deletions
Large diffs are not rendered by default.

‎stringdtype/stringdtype/src/dtype.c

Lines changed: 77 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -11,26 +11,35 @@ PyTypeObject *StringScalar_Type = NULL;
1111
PyObject *
1212
new_stringdtype_instance(PyObject *na_object, int coerce)
1313
{
14+
npy_string_allocator *allocator = NULL;
15+
PyThread_type_lock *allocator_lock = NULL;
16+
npy_packed_static_string packed_na_name = *NPY_EMPTY_STRING;
17+
npy_packed_static_string packed_default_string = *NPY_EMPTY_STRING;
18+
1419
PyObject *new =
1520
PyArrayDescr_Type.tp_new((PyTypeObject *)&StringDType, NULL, NULL);
1621

1722
if (new == NULL) {
1823
return NULL;
1924
}
2025

21-
npy_string_allocator *allocator = npy_string_new_allocator(
22-
PyMem_RawMalloc, PyMem_RawFree, PyMem_RawRealloc);
26+
allocator = npy_string_new_allocator(PyMem_RawMalloc, PyMem_RawFree,
27+
PyMem_RawRealloc);
2328

2429
if (allocator == NULL) {
2530
PyErr_SetString(PyExc_MemoryError,
2631
"Failed to create string allocator");
27-
return NULL;
32+
goto fail;
33+
}
34+
35+
allocator_lock = PyThread_allocate_lock();
36+
if (allocator_lock == NULL) {
37+
PyErr_SetString(PyExc_MemoryError, "Unable to allocate thread lock");
38+
goto fail;
2839
}
2940

3041
Py_XINCREF(na_object);
3142
((StringDTypeObject *)new)->na_object = na_object;
32-
npy_packed_static_string packed_na_name = *NPY_EMPTY_STRING;
33-
npy_packed_static_string packed_default_string = *NPY_EMPTY_STRING;
3443
int hasnull = na_object != NULL;
3544
int has_nan_na = 0;
3645
int has_string_na = 0;
@@ -45,8 +54,7 @@ new_stringdtype_instance(PyObject *na_object, int coerce)
4554
PyErr_SetString(PyExc_MemoryError,
4655
"Failed to allocate string while creating "
4756
"StringDType instance.");
48-
Py_DECREF(new);
49-
return NULL;
57+
goto fail;
5058
}
5159
}
5260
else {
@@ -55,8 +63,7 @@ new_stringdtype_instance(PyObject *na_object, int coerce)
5563
// NaN-like object)
5664
PyObject *eq = PyObject_RichCompare(na_object, na_object, Py_NE);
5765
if (eq == NULL) {
58-
Py_DECREF(new);
59-
return NULL;
66+
goto fail;
6067
}
6168
int is_truthy = PyObject_IsTrue(na_object);
6269
if (is_truthy == -1) {
@@ -70,25 +77,22 @@ new_stringdtype_instance(PyObject *na_object, int coerce)
7077
}
7178
PyObject *na_pystr = PyObject_Str(na_object);
7279
if (na_pystr == NULL) {
73-
Py_DECREF(new);
74-
return NULL;
80+
goto fail;
7581
}
7682

7783
Py_ssize_t size = 0;
7884
const char *utf8_ptr = PyUnicode_AsUTF8AndSize(na_pystr, &size);
7985
if (utf8_ptr == NULL) {
80-
Py_DECREF(new);
8186
Py_DECREF(na_pystr);
82-
return NULL;
87+
goto fail;
8388
}
8489
if (npy_string_newsize(utf8_ptr, (size_t)size, &packed_na_name,
8590
allocator) < 0) {
8691
PyErr_SetString(PyExc_MemoryError,
8792
"Failed to allocate string while creating "
8893
"StringDType instance.");
89-
Py_DECREF(new);
9094
Py_DECREF(na_pystr);
91-
return NULL;
95+
goto fail;
9296
}
9397
Py_DECREF(na_pystr);
9498
}
@@ -100,6 +104,7 @@ new_stringdtype_instance(PyObject *na_object, int coerce)
100104
snew->packed_default_string = packed_default_string;
101105
snew->packed_na_name = packed_na_name;
102106
snew->coerce = coerce;
107+
snew->allocator_lock = allocator_lock;
103108
snew->allocator = allocator;
104109
snew->array_owned = 0;
105110

@@ -140,6 +145,19 @@ new_stringdtype_instance(PyObject *na_object, int coerce)
140145
}
141146

142147
return new;
148+
149+
fail:
150+
// this only makes sense if the allocator isn't attached to new yet
151+
Py_DECREF(new);
152+
if (allocator != NULL) {
153+
npy_string_free(&packed_na_name, allocator);
154+
npy_string_free(&packed_default_string, allocator);
155+
npy_string_free_allocator(allocator);
156+
}
157+
if (allocator_lock != NULL) {
158+
PyThread_free_lock(allocator_lock);
159+
}
160+
return NULL;
143161
}
144162

145163
// sets the logical rules for determining equality between dtype instances
@@ -268,12 +286,14 @@ stringdtype_setitem(StringDTypeObject *descr, PyObject *obj, char **dataptr)
268286
{
269287
npy_packed_static_string *sdata = (npy_packed_static_string *)dataptr;
270288

289+
NPY_STRING_ACQUIRE_ALLOCATOR(descr);
290+
271291
// free if dataptr holds preexisting string data,
272292
// npy_string_free does a NULL check and checks for small strings
273293
if (npy_string_free(sdata, descr->allocator) < 0) {
274294
PyErr_SetString(PyExc_MemoryError,
275295
"String deallocation failed in StringDType setitem");
276-
return -1;
296+
goto fail;
277297
}
278298

279299
// borrow reference
@@ -288,26 +308,33 @@ stringdtype_setitem(StringDTypeObject *descr, PyObject *obj, char **dataptr)
288308
PyObject *val_obj = get_value(obj, descr->coerce);
289309

290310
if (val_obj == NULL) {
291-
return -1;
311+
goto fail;
292312
}
293313

294314
char *val = NULL;
295315
Py_ssize_t length = 0;
296316
if (PyBytes_AsStringAndSize(val_obj, &val, &length) == -1) {
297317
Py_DECREF(val_obj);
298-
return -1;
318+
goto fail;
299319
}
300320

301321
if (npy_string_newsize(val, length, sdata, descr->allocator) < 0) {
302322
PyErr_SetString(PyExc_MemoryError,
303323
"Failed to allocate string during StringDType "
304324
"setitem");
305325
Py_DECREF(val_obj);
306-
return -1;
326+
goto fail;
307327
}
308328
}
309329

330+
NPY_STRING_RELEASE_ALLOCATOR(descr);
331+
310332
return 0;
333+
334+
fail:
335+
NPY_STRING_RELEASE_ALLOCATOR(descr);
336+
337+
return -1;
311338
}
312339

313340
static PyObject *
@@ -317,12 +344,13 @@ stringdtype_getitem(StringDTypeObject *descr, char **dataptr)
317344
npy_packed_static_string *psdata = (npy_packed_static_string *)dataptr;
318345
npy_static_string sdata = {0, NULL};
319346
int hasnull = descr->na_object != NULL;
347+
NPY_STRING_ACQUIRE_ALLOCATOR(descr);
320348
int is_null = npy_string_load(descr->allocator, psdata, &sdata);
321349

322350
if (is_null < 0) {
323351
PyErr_SetString(PyExc_MemoryError,
324352
"Failed to load string in StringDType getitem");
325-
return NULL;
353+
goto fail;
326354
}
327355
else if (is_null == 1) {
328356
if (hasnull) {
@@ -331,16 +359,19 @@ stringdtype_getitem(StringDTypeObject *descr, char **dataptr)
331359
val_obj = na_object;
332360
}
333361
else {
362+
// cannot fail
334363
val_obj = PyUnicode_FromStringAndSize("", 0);
335364
}
336365
}
337366
else {
338367
val_obj = PyUnicode_FromStringAndSize(sdata.buf, sdata.size);
339368
if (val_obj == NULL) {
340-
return NULL;
369+
goto fail;
341370
}
342371
}
343372

373+
NPY_STRING_RELEASE_ALLOCATOR(descr);
374+
344375
/*
345376
* In principle we should return a StringScalar instance here, but
346377
* creating a StringScalar via PyObject_CallFunctionObjArgs has
@@ -353,6 +384,12 @@ stringdtype_getitem(StringDTypeObject *descr, char **dataptr)
353384
* eventually the scalar type for this dtype will be str.
354385
*/
355386
return val_obj;
387+
388+
fail:
389+
390+
NPY_STRING_RELEASE_ALLOCATOR(descr);
391+
392+
return NULL;
356393
}
357394

358395
// PyArray_NonzeroFunc
@@ -369,7 +406,10 @@ int
369406
compare(void *a, void *b, void *arr)
370407
{
371408
StringDTypeObject *descr = (StringDTypeObject *)PyArray_DESCR(arr);
372-
return _compare(a, b, descr, descr);
409+
NPY_STRING_ACQUIRE_ALLOCATOR(descr);
410+
int ret = _compare(a, b, descr, descr);
411+
NPY_STRING_RELEASE_ALLOCATOR(descr);
412+
return ret;
373413
}
374414

375415
int
@@ -480,20 +520,25 @@ stringdtype_clear_loop(void *NPY_UNUSED(traverse_context),
480520
npy_intp stride, NpyAuxData *NPY_UNUSED(auxdata))
481521
{
482522
StringDTypeObject *sdescr = (StringDTypeObject *)descr;
523+
NPY_STRING_ACQUIRE_ALLOCATOR(sdescr);
483524
while (size--) {
484525
npy_packed_static_string *sdata = (npy_packed_static_string *)data;
485526
if (data != NULL) {
486527
if (npy_string_free(sdata, sdescr->allocator) < 0) {
487528
gil_error(PyExc_MemoryError,
488529
"String deallocation failed in clear loop");
489-
return -1;
530+
goto fail;
490531
}
491532
memset(data, 0, sizeof(npy_packed_static_string));
492533
}
493534
data += stride;
494535
}
495-
536+
NPY_STRING_RELEASE_ALLOCATOR(sdescr);
496537
return 0;
538+
539+
fail:
540+
NPY_STRING_RELEASE_ALLOCATOR(sdescr);
541+
return -1;
497542
}
498543

499544
static int
@@ -652,9 +697,14 @@ static void
652697
stringdtype_dealloc(StringDTypeObject *self)
653698
{
654699
Py_XDECREF(self->na_object);
655-
npy_string_free(&self->packed_default_string, self->allocator);
656-
npy_string_free(&self->packed_na_name, self->allocator);
657-
npy_string_free_allocator(self->allocator);
700+
// this can be null if an error happens while initializing an instance
701+
if (self->allocator != NULL) {
702+
// can we assume the destructor for an instance will only get called
703+
// inside of one C thread?
704+
npy_string_free(&self->packed_default_string, self->allocator);
705+
npy_string_free(&self->packed_na_name, self->allocator);
706+
npy_string_free_allocator(self->allocator);
707+
}
658708
PyArrayDescr_Type.tp_dealloc((PyObject *)self);
659709
}
660710

‎stringdtype/stringdtype/src/dtype.h

Lines changed: 44 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,42 @@
1919
#include "numpy/npy_math.h"
2020
#include "numpy/ufuncobject.h"
2121

22+
#define NPY_STRING_ACQUIRE_ALLOCATOR(descr) \
23+
if (!PyThread_acquire_lock(descr->allocator_lock, NOWAIT_LOCK)) { \
24+
PyThread_acquire_lock(descr->allocator_lock, WAIT_LOCK); \
25+
}
26+
27+
#define NPY_STRING_ACQUIRE_ALLOCATOR2(descr1, descr2) \
28+
NPY_STRING_ACQUIRE_ALLOCATOR(descr1) \
29+
if (descr1 != descr2) { \
30+
NPY_STRING_ACQUIRE_ALLOCATOR(descr2) \
31+
}
32+
33+
#define NPY_STRING_ACQUIRE_ALLOCATOR3(descr1, descr2, descr3) \
34+
NPY_STRING_ACQUIRE_ALLOCATOR(descr1) \
35+
if (descr1 != descr2) { \
36+
NPY_STRING_ACQUIRE_ALLOCATOR(descr2) \
37+
} \
38+
if (descr1 != descr3 && descr2 != descr3) { \
39+
NPY_STRING_ACQUIRE_ALLOCATOR(descr3) \
40+
}
41+
42+
#define NPY_STRING_RELEASE_ALLOCATOR(descr) \
43+
PyThread_release_lock(descr->allocator_lock);
44+
#define NPY_STRING_RELEASE_ALLOCATOR2(descr1, descr2) \
45+
NPY_STRING_RELEASE_ALLOCATOR(descr1); \
46+
if (descr1 != descr2) { \
47+
NPY_STRING_RELEASE_ALLOCATOR(descr2); \
48+
}
49+
#define NPY_STRING_RELEASE_ALLOCATOR3(descr1, descr2, descr3) \
50+
NPY_STRING_RELEASE_ALLOCATOR(descr1); \
51+
if (descr1 != descr2) { \
52+
NPY_STRING_RELEASE_ALLOCATOR(descr2); \
53+
} \
54+
if (descr1 != descr3 && descr2 != descr3) { \
55+
NPY_STRING_RELEASE_ALLOCATOR(descr3); \
56+
}
57+
2258
typedef struct {
2359
PyArray_Descr base;
2460
PyObject *na_object;
@@ -30,6 +66,11 @@ typedef struct {
3066
npy_packed_static_string packed_default_string;
3167
npy_static_string na_name;
3268
npy_packed_static_string packed_na_name;
69+
PyThread_type_lock *allocator_lock;
70+
// the allocator should only be directly accessed after
71+
// acquiring the allocator_lock and the lock should
72+
// be released immediately after the allocator is
73+
// no longer needed
3374
npy_string_allocator *allocator;
3475
} StringDTypeObject;
3576

@@ -46,6 +87,8 @@ new_stringdtype_instance(PyObject *na_object, int coerce);
4687
int
4788
init_string_dtype(void);
4889

90+
// Assumes that the caller has already acquired the allocator locks for both
91+
// descriptors
4992
int
5093
_compare(void *a, void *b, StringDTypeObject *descr_a,
5194
StringDTypeObject *descr_b);
@@ -60,9 +103,7 @@ stringdtype_setitem(StringDTypeObject *descr, PyObject *obj, char **dataptr);
60103
void
61104
gil_error(PyObject *type, const char *msg);
62105

63-
// from dtypemeta.h, not public in numpy
64-
#define NPY_DTYPE(descr) ((PyArray_DTypeMeta *)Py_TYPE(descr))
65-
106+
// the locks on both allocators must be acquired before calling this function
66107
int
67108
free_and_copy(npy_string_allocator *in_allocator,
68109
npy_string_allocator *out_allocator,

0 commit comments

Comments
 (0)