Added memory leak fix and testing for Encoder

This commit is contained in:
Josh Ware 2014-08-27 15:30:31 +10:00 committed by wiredfool
parent bb738aef38
commit caa95a26b2
4 changed files with 39 additions and 1 deletions

View File

@ -473,6 +473,7 @@ def _save(im, fp, tile, bufsize=0):
break break
if s < 0: if s < 0:
raise IOError("encoder error %d when writing image file" % s) raise IOError("encoder error %d when writing image file" % s)
e.cleanup()
else: else:
# slight speedup: compress to real file object # slight speedup: compress to real file object
for e, b, o, a in tile: for e, b, o, a in tile:
@ -483,6 +484,7 @@ def _save(im, fp, tile, bufsize=0):
s = e.encode_to_file(fh, bufsize) s = e.encode_to_file(fh, bufsize)
if s < 0: if s < 0:
raise IOError("encoder error %d when writing image file" % s) raise IOError("encoder error %d when writing image file" % s)
e.cleanup()
try: try:
fp.flush() fp.flush()
except: pass except: pass

View File

@ -8,6 +8,8 @@ stack_size = 8*1048576
iterations = int((mem_limit/stack_size)*2) iterations = int((mem_limit/stack_size)*2)
codecs = dir(Image.core) codecs = dir(Image.core)
test_file = "Tests/images/rgb_trns_ycbc.jp2" test_file = "Tests/images/rgb_trns_ycbc.jp2"
from commands import getoutput
from os import getpid
@unittest.skipIf(sys.platform.startswith('win32'), "requires Unix or MacOS") @unittest.skipIf(sys.platform.startswith('win32'), "requires Unix or MacOS")
@ -16,7 +18,7 @@ class TestJpegLeaks(PillowTestCase):
if "jpeg2k_encoder" not in codecs or "jpeg2k_decoder" not in codecs: if "jpeg2k_encoder" not in codecs or "jpeg2k_decoder" not in codecs:
self.skipTest('JPEG 2000 support not available') self.skipTest('JPEG 2000 support not available')
def test_leak(self): def test_leak_load(self):
from resource import setrlimit, RLIMIT_AS, RLIMIT_STACK from resource import setrlimit, RLIMIT_AS, RLIMIT_STACK
setrlimit(RLIMIT_STACK, (stack_size, stack_size)) setrlimit(RLIMIT_STACK, (stack_size, stack_size))
setrlimit(RLIMIT_AS, (mem_limit, mem_limit)) setrlimit(RLIMIT_AS, (mem_limit, mem_limit))
@ -24,5 +26,22 @@ class TestJpegLeaks(PillowTestCase):
with Image.open(test_file) as im: with Image.open(test_file) as im:
im.load() im.load()
def test_leak_save(self):
from resource import setrlimit, RLIMIT_AS, RLIMIT_STACK
try:
from cStringIO import StringIO
except ImportError:
from io import StringIO
setrlimit(RLIMIT_STACK, (stack_size, stack_size))
setrlimit(RLIMIT_AS, (mem_limit, mem_limit))
for count in range(iterations):
with Image.open(test_file) as im:
im.load()
test_output = StringIO()
im.save(test_output, "JPEG2000")
test_output.seek(0)
output = test_output.read()
if __name__ == '__main__': if __name__ == '__main__':
unittest.main() unittest.main()

View File

@ -99,6 +99,18 @@ _dealloc(ImagingEncoderObject* encoder)
PyObject_Del(encoder); PyObject_Del(encoder);
} }
static PyObject*
_encode_cleanup(ImagingEncoderObject* encoder, PyObject* args)
{
int status = 0;
if (encoder->cleanup){
status = encoder->cleanup(&encoder->state);
}
return Py_BuildValue("i", status);
}
static PyObject* static PyObject*
_encode(ImagingEncoderObject* encoder, PyObject* args) _encode(ImagingEncoderObject* encoder, PyObject* args)
{ {
@ -239,6 +251,7 @@ _setimage(ImagingEncoderObject* encoder, PyObject* args)
static struct PyMethodDef methods[] = { static struct PyMethodDef methods[] = {
{"encode", (PyCFunction)_encode, 1}, {"encode", (PyCFunction)_encode, 1},
{"cleanup", (PyCFunction)_encode_cleanup, 1},
{"encode_to_file", (PyCFunction)_encode_to_file, 1}, {"encode_to_file", (PyCFunction)_encode_to_file, 1},
{"setimage", (PyCFunction)_setimage, 1}, {"setimage", (PyCFunction)_setimage, 1},
{NULL, NULL} /* sentinel */ {NULL, NULL} /* sentinel */

View File

@ -585,6 +585,10 @@ ImagingJpeg2KEncodeCleanup(ImagingCodecState state) {
if (context->encoder) if (context->encoder)
ImagingIncrementalCodecDestroy(context->encoder); ImagingIncrementalCodecDestroy(context->encoder);
/* Prevent multiple calls to ImagingIncrementalCodecDestroy */
context->encoder = NULL;
return -1; return -1;
} }