Skip to content

bpo-28416: Break reference cycles in Pickler and Unpickler subclasses #4080

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
65 changes: 63 additions & 2 deletions Lib/test/test_pickle.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import collections
import struct
import sys
import weakref

import unittest
from test import support
Expand Down Expand Up @@ -112,6 +113,66 @@ class PyIdPersPicklerTests(AbstractIdentityPersistentPicklerTests,
pickler = pickle._Pickler
unpickler = pickle._Unpickler

@support.cpython_only
def test_pickler_reference_cycle(self):
def check(Pickler):
for bin in 0, 1:
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmm, what is bin? Do you mean protocol?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

A flag text/binary. Later it became a protocol, but the attribute bin still is used in sources (it is equal protocol != 0). Here it is enough to test only one binary protocol.

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes... It's still a bit cryptic though. And it doesn't cost anything to loop over all protocols, right? :-)

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done.

f = io.BytesIO()
pickler = Pickler(f, bin)
pickler.dump('abc')
self.assertEqual(self.loads(f.getvalue()), 'abc')
pickler = Pickler(io.BytesIO())
self.assertEqual(pickler.persistent_id('def'), 'def')
r = weakref.ref(pickler)
del pickler
self.assertIsNone(r())

class PersPickler(self.pickler):
def persistent_id(subself, obj):
return obj
check(PersPickler)

class PersPickler(self.pickler):
@classmethod
def persistent_id(cls, obj):
return obj
check(PersPickler)

class PersPickler(self.pickler):
@staticmethod
def persistent_id(obj):
return obj
check(PersPickler)

@support.cpython_only
def test_unpickler_reference_cycle(self):
def check(Unpickler):
for bin in 0, 1:
unpickler = Unpickler(io.BytesIO(self.dumps('abc', bin)))
self.assertEqual(unpickler.load(), 'abc')
unpickler = Unpickler(io.BytesIO())
self.assertEqual(unpickler.persistent_load('def'), 'def')
r = weakref.ref(unpickler)
del unpickler
self.assertIsNone(r())

class PersUnpickler(self.unpickler):
def persistent_load(subself, pid):
return pid
check(PersUnpickler)

class PersUnpickler(self.unpickler):
@classmethod
def persistent_load(cls, pid):
return pid
check(PersUnpickler)

class PersUnpickler(self.unpickler):
@staticmethod
def persistent_load(pid):
return pid
check(PersUnpickler)


class PyPicklerUnpicklerObjectTests(AbstractPicklerUnpicklerObjectTests):

Expand Down Expand Up @@ -189,7 +250,7 @@ class SizeofTests(unittest.TestCase):
check_sizeof = support.check_sizeof

def test_pickler(self):
basesize = support.calcobjsize('5P2n3i2n3iP')
basesize = support.calcobjsize('2Pi3P2n3i2n3iP')
p = _pickle.Pickler(io.BytesIO())
self.assertEqual(object.__sizeof__(p), basesize)
MT_size = struct.calcsize('3nP0n')
Expand All @@ -206,7 +267,7 @@ def test_pickler(self):
0) # Write buffer is cleared after every dump().

def test_unpickler(self):
basesize = support.calcobjsize('2Pn2P 2P2n2i5P 2P3n6P2n2i')
basesize = support.calcobjsize('2P2nPi 2P2n2i5P 2P3n6P2n2i')
unpickler = _pickle.Unpickler
P = struct.calcsize('P') # Size of memo table entry.
n = struct.calcsize('n') # Size of mark table entry.
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,3 @@
Instances of pickle.Pickler subclass with the persistent_id() method and
pickle.Unpickler subclass with the persistent_load() method no longer create
reference cycles.
153 changes: 115 additions & 38 deletions Modules/_pickle.c
Original file line number Diff line number Diff line change
Expand Up @@ -360,6 +360,81 @@ _Pickle_FastCall(PyObject *func, PyObject *obj)

/*************************************************************************/

static int
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I wonder if the role of those three functions would be clearer if they took a {PyObject *func, int unbound} structure.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't know. This function returns a pair: a pointer to bound or unbound method, and a flag for distinguishing them. There are other functions that return the same pair: _PyObject_GetMethod() and lookup_maybe_method() in typeobject.c. But the interface is different. _PyObject_GetMethod() returns a flag and sets a pointer by reference, lookup_maybe_method() returns a pointer and sets a flag by reference, this functions sets both pointer and flag by reference (and it also decrement the refcount of the old pointer). I tried different variants, but I don't know what interface is better. Well, will try to use a structure.

init_method_ref(PyObject *self, _Py_Identifier *name,
PyObject **method, int *unbound)
{
PyObject *func, *func2;

/* *method and *unbound should be consistent. All refcount decrements
should be occurred after setting *unbound and *method. */
*unbound = 0;
Py_CLEAR(*method);

func = _PyObject_GetAttrId(self, name);
if (func == NULL) {
*unbound = 0;
Py_CLEAR(*method);
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The two lines above seem redundant with the function beginning.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actually lines at function beginning are redundant.

if (!PyErr_ExceptionMatches(PyExc_AttributeError)) {
return -1;
}
PyErr_Clear();
return 0;
}

/* Avoid creating a reference cycle (pickler -> bound method of pickler ->
-> pickler). */
if (PyMethod_Check(func) && PyMethod_GET_SELF(func) == self) {
/* bound Python method */
func2 = PyMethod_GET_FUNCTION(func);
Py_INCREF(func2);
}
else if (PyCFunction_Check(func) && PyCFunction_GET_SELF(func) == self) {
/* bound built-in method */
func2 = PyDescr_NewMethod(Py_TYPE(self),
((PyCFunctionObject *)func)->m_ml);
if (func2 == NULL) {
Py_DECREF(func);
return -1;
}
}
else {
*unbound = 0;
Py_XSETREF(*method, func);
return 0;
}
*unbound = 1;
Py_XSETREF(*method, func2);
Py_DECREF(func);
return 0;
}

static PyObject *
bound_method(PyObject *self, PyObject *func, int unbound)
{
if (unbound) {
descrgetfunc f = func->ob_type->tp_descr_get;
if (f != NULL) {
return f(func, self, (PyObject *)Py_TYPE(self));
}
}
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Why does this fall through? May we return a non-bound method when unbound is true?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmm, interesting question.

Py_INCREF(func);
return func;
}

static PyObject *
call_method(PyObject *self, PyObject *func, int unbound, PyObject *obj)
{
if (unbound) {
return PyObject_CallFunctionObjArgs(func, self, obj, NULL);
}
else {
return PyObject_CallFunctionObjArgs(func, obj, NULL);
}
}

/*************************************************************************/

/* Internal data type used as the unpickling stack. */
typedef struct {
PyObject_VAR_HEAD
Expand Down Expand Up @@ -552,6 +627,8 @@ typedef struct PicklerObject {
objects to support self-referential objects
pickling. */
PyObject *pers_func; /* persistent_id() method, can be NULL */
int unbound_pers_func; /* 1 if pers_func is unbound method,
0 if pers_func is bound method */
PyObject *dispatch_table; /* private dispatch_table, can be NULL */

PyObject *write; /* write() method of the output stream. */
Expand Down Expand Up @@ -590,6 +667,8 @@ typedef struct UnpicklerObject {
Py_ssize_t memo_len; /* Number of objects in the memo */

PyObject *pers_func; /* persistent_load() method, can be NULL. */
int unbound_pers_func; /* 1 if pers_func is unbound method,
0 if pers_func is bound method */

Py_buffer buffer;
char *input_buffer;
Expand Down Expand Up @@ -3440,16 +3519,16 @@ save_type(PicklerObject *self, PyObject *obj)
}

static int
save_pers(PicklerObject *self, PyObject *obj, PyObject *func)
save_pers(PicklerObject *self, PyObject *obj)
{
PyObject *pid = NULL;
int status = 0;

const char persid_op = PERSID;
const char binpersid_op = BINPERSID;

Py_INCREF(obj);
pid = _Pickle_FastCall(func, obj);
pid = call_method((PyObject *)self,
self->pers_func, self->unbound_pers_func, obj);
if (pid == NULL)
return -1;

Expand Down Expand Up @@ -3827,7 +3906,7 @@ save(PicklerObject *self, PyObject *obj, int pers_save)
0 if it did nothing successfully;
1 if a persistent id was saved.
*/
if ((status = save_pers(self, obj, self->pers_func)) != 0)
if ((status = save_pers(self, obj)) != 0)
goto done;
}

Expand Down Expand Up @@ -4242,13 +4321,10 @@ _pickle_Pickler___init___impl(PicklerObject *self, PyObject *file,
self->fast_nesting = 0;
self->fast_memo = NULL;

self->pers_func = _PyObject_GetAttrId((PyObject *)self,
&PyId_persistent_id);
if (self->pers_func == NULL) {
if (!PyErr_ExceptionMatches(PyExc_AttributeError)) {
return -1;
}
PyErr_Clear();
if (init_method_ref((PyObject *)self, &PyId_persistent_id,
&self->pers_func, &self->unbound_pers_func) < 0)
{
return -1;
}

self->dispatch_table = _PyObject_GetAttrId((PyObject *)self,
Expand Down Expand Up @@ -4515,11 +4591,12 @@ Pickler_set_memo(PicklerObject *self, PyObject *obj)
static PyObject *
Pickler_get_persid(PicklerObject *self)
{
if (self->pers_func == NULL)
if (self->pers_func == NULL) {
PyErr_SetString(PyExc_AttributeError, "persistent_id");
else
Py_INCREF(self->pers_func);
return self->pers_func;
return NULL;
}
return bound_method((PyObject *)self,
self->pers_func, self->unbound_pers_func);
}

static int
Expand All @@ -4536,6 +4613,7 @@ Pickler_set_persid(PicklerObject *self, PyObject *value)
return -1;
}

self->unbound_pers_func = 0;
Py_INCREF(value);
Py_XSETREF(self->pers_func, value);

Expand Down Expand Up @@ -5485,7 +5563,7 @@ load_stack_global(UnpicklerObject *self)
static int
load_persid(UnpicklerObject *self)
{
PyObject *pid;
PyObject *pid, *obj;
Py_ssize_t len;
char *s;

Expand All @@ -5505,13 +5583,13 @@ load_persid(UnpicklerObject *self)
return -1;
}

/* This does not leak since _Pickle_FastCall() steals the reference
to pid first. */
pid = _Pickle_FastCall(self->pers_func, pid);
if (pid == NULL)
obj = call_method((PyObject *)self,
self->pers_func, self->unbound_pers_func, pid);
Py_DECREF(pid);
if (obj == NULL)
return -1;

PDATA_PUSH(self->stack, pid, -1);
PDATA_PUSH(self->stack, obj, -1);
return 0;
}
else {
Expand All @@ -5526,20 +5604,20 @@ load_persid(UnpicklerObject *self)
static int
load_binpersid(UnpicklerObject *self)
{
PyObject *pid;
PyObject *pid, *obj;

if (self->pers_func) {
PDATA_POP(self->stack, pid);
if (pid == NULL)
return -1;

/* This does not leak since _Pickle_FastCall() steals the
reference to pid first. */
pid = _Pickle_FastCall(self->pers_func, pid);
if (pid == NULL)
obj = call_method((PyObject *)self,
self->pers_func, self->unbound_pers_func, pid);
Py_DECREF(pid);
if (obj == NULL)
return -1;

PDATA_PUSH(self->stack, pid, -1);
PDATA_PUSH(self->stack, obj, -1);
return 0;
}
else {
Expand Down Expand Up @@ -6686,13 +6764,10 @@ _pickle_Unpickler___init___impl(UnpicklerObject *self, PyObject *file,

self->fix_imports = fix_imports;

self->pers_func = _PyObject_GetAttrId((PyObject *)self,
&PyId_persistent_load);
if (self->pers_func == NULL) {
if (!PyErr_ExceptionMatches(PyExc_AttributeError)) {
return -1;
}
PyErr_Clear();
if (init_method_ref((PyObject *)self, &PyId_persistent_load,
&self->pers_func, &self->unbound_pers_func) < 0)
{
return -1;
}

self->stack = (Pdata *)Pdata_New();
Expand Down Expand Up @@ -6979,11 +7054,12 @@ Unpickler_set_memo(UnpicklerObject *self, PyObject *obj)
static PyObject *
Unpickler_get_persload(UnpicklerObject *self)
{
if (self->pers_func == NULL)
if (self->pers_func == NULL) {
PyErr_SetString(PyExc_AttributeError, "persistent_load");
else
Py_INCREF(self->pers_func);
return self->pers_func;
return NULL;
}
return bound_method((PyObject *)self,
self->pers_func, self->unbound_pers_func);
}

static int
Expand All @@ -7001,6 +7077,7 @@ Unpickler_set_persload(UnpicklerObject *self, PyObject *value)
return -1;
}

self->unbound_pers_func = 0;
Py_INCREF(value);
Py_XSETREF(self->pers_func, value);

Expand Down