Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 499ab6a

Browse files
jeremyhyltonjeremyhylton
committed
Better fix for core dumps on recursive objects in fast mode.
Raise ValueError when an object contains an arbitrarily nested reference to itself. (The previous fix just produced invalid pickles.) Solution is very much like Py_ReprEnter() and Py_ReprLeave(): fast_save_enter() and fast_save_leave() that tracks the fast_container limit and keeps a fast_memo of objects currently being pickled. The cost of the solution is moderately expensive for deeply nested structures, but it still seems to be faster than normal pickling, based on tests with deeply nested lists. Once FAST_LIMIT is exceeded, the new code is about twice as slow as fast-mode code that doesn't check for recursion. It's still twice as fast as the normal pickling code. In the absence of deeply nested structures, I couldn't measure a difference.
1 parent abe2c62 commit 499ab6a

1 file changed

Lines changed: 64 additions & 32 deletions

File tree

‎Modules/cPickle.c‎

Lines changed: 64 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -318,6 +318,7 @@ typedef struct Picklerobject {
318318
int buf_size;
319319
PyObject *dispatch_table;
320320
int fast_container; /* count nested container dumps */
321+
PyObject *fast_memo;
321322
} Picklerobject;
322323

323324
#define FAST_LIMIT 2000
@@ -886,6 +887,51 @@ whichmodule(PyObject *global, PyObject *global_name) {
886887
}
887888

888889

890+
static int
891+
fast_save_enter(Picklerobject *self, PyObject *obj)
892+
{
893+
/* if fast_container < 0, we're doing an error exit. */
894+
if (++self->fast_container >= FAST_LIMIT) {
895+
PyObject *key = NULL;
896+
if (self->fast_memo == NULL) {
897+
self->fast_memo = PyDict_New();
898+
if (self->fast_memo == NULL) {
899+
self->fast_container = -1;
900+
return 0;
901+
}
902+
}
903+
key = PyLong_FromVoidPtr(obj);
904+
if (key == NULL)
905+
return 0;
906+
if (PyDict_GetItem(self->fast_memo, key)) {
907+
PyErr_Format(PyExc_ValueError,
908+
"fast mode: can't pickle cyclic objects including object type %s at %p",
909+
obj->ob_type->tp_name, obj);
910+
self->fast_container = -1;
911+
return 0;
912+
}
913+
if (PyDict_SetItem(self->fast_memo, key, Py_None) < 0) {
914+
self->fast_container = -1;
915+
return 0;
916+
}
917+
}
918+
return 1;
919+
}
920+
921+
int
922+
fast_save_leave(Picklerobject *self, PyObject *obj)
923+
{
924+
if (self->fast_container-- >= FAST_LIMIT) {
925+
PyObject *key = PyLong_FromVoidPtr(obj);
926+
if (key == NULL)
927+
return 0;
928+
if (PyDict_DelItem(self->fast_memo, key) < 0) {
929+
return 0;
930+
}
931+
}
932+
return 1;
933+
}
934+
889935
static int
890936
save_none(Picklerobject *self, PyObject *args) {
891937
static char none = NONE;
@@ -1357,15 +1403,13 @@ save_empty_tuple(Picklerobject *self, PyObject *args) {
13571403
static int
13581404
save_list(Picklerobject *self, PyObject *args) {
13591405
PyObject *element = 0;
1360-
int s_len, len, i, using_appends, res = -1, unfast = 0;
1406+
int s_len, len, i, using_appends, res = -1;
13611407
char s[3];
13621408

13631409
static char append = APPEND, appends = APPENDS;
13641410

1365-
if (self->fast && self->fast_container++ > FAST_LIMIT) {
1366-
self->fast = 0;
1367-
unfast = 1;
1368-
}
1411+
if (self->fast && !fast_save_enter(self, args))
1412+
goto finally;
13691413

13701414
if (self->bin) {
13711415
s[0] = EMPTY_LIST;
@@ -1417,11 +1461,8 @@ save_list(Picklerobject *self, PyObject *args) {
14171461
res = 0;
14181462

14191463
finally:
1420-
if (self->fast || unfast) {
1421-
self->fast_container--;
1422-
if (unfast && self->fast_container < FAST_LIMIT)
1423-
self->fast = 1;
1424-
}
1464+
if (self->fast && !fast_save_leave(self, args))
1465+
res = -1;
14251466

14261467
return res;
14271468
}
@@ -1430,15 +1471,13 @@ save_list(Picklerobject *self, PyObject *args) {
14301471
static int
14311472
save_dict(Picklerobject *self, PyObject *args) {
14321473
PyObject *key = 0, *value = 0;
1433-
int i, len, res = -1, using_setitems, unfast = 0;
1474+
int i, len, res = -1, using_setitems;
14341475
char s[3];
14351476

14361477
static char setitem = SETITEM, setitems = SETITEMS;
14371478

1438-
if (self->fast && self->fast_container++ > FAST_LIMIT) {
1439-
self->fast = 0;
1440-
unfast = 1;
1441-
}
1479+
if (self->fast && !fast_save_enter(self, args))
1480+
goto finally;
14421481

14431482
if (self->bin) {
14441483
s[0] = EMPTY_DICT;
@@ -1491,12 +1530,8 @@ save_dict(Picklerobject *self, PyObject *args) {
14911530
res = 0;
14921531

14931532
finally:
1494-
if (self->fast || unfast) {
1495-
self->fast_container--;
1496-
if (unfast && self->fast_container < FAST_LIMIT)
1497-
self->fast = 1;
1498-
}
1499-
1533+
if (self->fast && !fast_save_leave(self, args))
1534+
res = -1;
15001535

15011536
return res;
15021537
}
@@ -1507,14 +1542,12 @@ save_inst(Picklerobject *self, PyObject *args) {
15071542
PyObject *class = 0, *module = 0, *name = 0, *state = 0,
15081543
*getinitargs_func = 0, *getstate_func = 0, *class_args = 0;
15091544
char *module_str, *name_str;
1510-
int module_size, name_size, res = -1, unfast = 0;
1545+
int module_size, name_size, res = -1;
15111546

15121547
static char inst = INST, obj = OBJ, build = BUILD;
15131548

1514-
if (self->fast && self->fast_container++ > FAST_LIMIT) {
1515-
self->fast = 0;
1516-
unfast = 1;
1517-
}
1549+
if (self->fast && !fast_save_enter(self, args))
1550+
goto finally;
15181551

15191552
if ((*self->write_func)(self, &MARKv, 1) < 0)
15201553
goto finally;
@@ -1622,11 +1655,8 @@ save_inst(Picklerobject *self, PyObject *args) {
16221655
res = 0;
16231656

16241657
finally:
1625-
if (self->fast || unfast) {
1626-
self->fast_container--;
1627-
if (unfast && self->fast_container < FAST_LIMIT)
1628-
self->fast = 1;
1629-
}
1658+
if (self->fast && !fast_save_leave(self, args))
1659+
res = -1;
16301660

16311661
Py_XDECREF(module);
16321662
Py_XDECREF(class);
@@ -1669,7 +1699,7 @@ save_global(Picklerobject *self, PyObject *args, PyObject *name) {
16691699
mod = PyImport_ImportModule(module_str);
16701700
if (mod == NULL) {
16711701
/* Py_ErrClear(); ?? */
1672-
cPickle_ErrFormat(PicklingError,
1702+
cPickle_ErrFormat(PicklingError,
16731703
"Can't pickle %s: it's not found as %s.%s",
16741704
"OSS", args, module, global_name);
16751705
goto finally;
@@ -2251,6 +2281,7 @@ newPicklerobject(PyObject *file, int bin) {
22512281
self->bin = bin;
22522282
self->fast = 0;
22532283
self->fast_container = 0;
2284+
self->fast_memo = NULL;
22542285
self->buf_size = 0;
22552286
self->dispatch_table = NULL;
22562287

@@ -2339,6 +2370,7 @@ static void
23392370
Pickler_dealloc(Picklerobject *self) {
23402371
Py_XDECREF(self->write);
23412372
Py_XDECREF(self->memo);
2373+
Py_XDECREF(self->fast_memo);
23422374
Py_XDECREF(self->arg);
23432375
Py_XDECREF(self->file);
23442376
Py_XDECREF(self->pers_func);

0 commit comments

Comments
 (0)