emote collector memory leak debugging

emote_collector.utils.image.py

import contextlib
import imghdr
import io
import logging

logger = logging.getLogger(__name__)

try:
	import wand.image
except ImportError:
	logger.warn('Failed to import wand.image. Image manipulation functions will be unavailable.')

from . import asyncexecutor
from . import errors
from . import size

def _resize_until_small(image_data: io.BytesIO) -> io.BytesIO:
	"""If the image_data is bigger than 256KB, resize it until it's not.

	If resizing takes more than 30 seconds, raise asyncio.TimeoutError.
	"""
	# It's important that we only attempt to resize the image when we have to,
	# ie when it exceeds the Discord limit of 256KiB.
	# Apparently some <256KiB images become larger when we attempt to resize them,
	# so resizing sometimes does more harm than good.
	max_resolution = 128  # pixels
	image_size = size(image_data)
	while image_size > 256 * 2**10 and max_resolution >= 32:  # don't resize past 32x32 or 256KiB
		logger.debug('image size too big (%s bytes)', image_size)
		logger.debug('attempting resize to %s*%s pixels', max_resolution, max_resolution)
		image_data = thumbnail(image_data, (max_resolution, max_resolution))
		image_size = size(image_data)
		max_resolution //= 2
	return image_data

# allow testing this in the bot (using an executor), and standalone (blocking)
resize_until_small = asyncexecutor(timeout=30)(_resize_until_small)

def thumbnail(image_data: io.BytesIO, max_size=(128, 128)) -> io.BytesIO:
	"""Resize an image in place to no more than max_size pixels, preserving aspect ratio.

	Return the new image.
	"""
	# Credit to @Liara#0001 (ID 136900814408122368)
	# https://gitlab.com/Pandentia/element-zero/blob/47bc8eeeecc7d353ec66e1ef5235adab98ca9635/element_zero/cogs/emoji.py#L243-247
	with wand.image.Image(blob=image_data) as image:
		new_resolution = scale_resolution((image.width, image.height), max_size)
		image.resize(*new_resolution)
		# we create a new buffer here because there's wand errors otherwise.
		# specific error:
		# MissingDelegateError: no decode delegate for this image format `' @ error/blob.c/BlobToImage/353
		out = io.BytesIO()
		image.save(file=out)

	# allow resizing the original image more than once for memory profiling
	image_data.seek(0)
	# allow reading the resized image data
	out.seek(0)

	return out

def scale_resolution(old_res, new_res):
	"""Resize a resolution, preserving aspect ratio. Returned w,h will be <= new_res"""
	# https://stackoverflow.com/a/6565988

	old_width, old_height = old_res
	new_width, new_height = new_res

	old_ratio = old_width / old_height
	new_ratio = new_width / new_height

	if new_ratio > old_ratio:
		return (round(old_width * new_height/old_height), new_height)
	else:
		return (new_width, round(old_height * new_width/old_width))

def is_animated(image_data: bytes):
	"""Return whether the image data is animated, or raise InvalidImageError if it's not an image."""
	type = imghdr.what(None, image_data)
	if type == 'gif':
		return True
	elif type in {'png', 'jpeg'}:
		return False
	else:
		raise errors.InvalidImageError

memleak_test.py

#!/usr/bin/env python3
# encoding: utf-8

import io
import os.path

import objgraph

from ben_cogs.debug import Debug
from emote_collector.utils.image import _resize_until_small as resize_until_small

memory_usage = Debug().memory_usage

out = io.BytesIO()
with open(os.path.expanduser('~/Pictures/discord emoji/tests/big/terry davis.gif'), 'rb') as f:
	out.write(f.read())
out.seek(0)

def print_header(label):
	print()
	print(label)
	print('─' * len(label))

def show_memory_usage():
	print('Memory usage:', memory_usage())

def show_growth():
	show_memory_usage()
	print('Object growth:', end='\n\n')
	objgraph.show_growth(limit=0)

# collect a baseline
show_memory_usage()
objgraph.growth()

resized = resize_until_small(out)

print_header('after resize')
show_growth()

# do it again, to see if memory usage goes up a second time
# i think it may not, because the first time around, although mem usage went up,
# there was no abnormal growth in Image or BytesIO objects
resized = resize_until_small(out)

print_header('after 2nd resize')
show_growth()

for resize in range(1, 11):
	resized = resize_until_small(out)
	print_header(f'after {resize} additional resize(s)')
	show_growth()

# wait for me to check htop to see if memory usage goes back down
print()
input('Press enter to continue ')

# what's referencing these?
for image in objgraph.by_type('BytesIO'):
	print('found image')
	objgraph.show_chain(
		objgraph.find_backref_chain(
			image,
			objgraph.is_proper_module),
		filename=f'{image!r} chain.png')

output.txt

Memory usage: 32.8 MB

after resize
────────────
Memory usage: 295.4 MB
Object growth:

function                       9079       +23
tuple                          4850       +12
dict                           4577        +7
weakref                        2226        +6
set                             204        +5
property                       1099        +4
cell                            962        +3
builtin_function_or_method     1861        +2
module                          466        +1
type                           1215        +1
ModuleSpec                      461        +1
frozenset                       170        +1
SourceFileLoader                388        +1
ABCMeta                         134        +1
BytesIO                           2        +1

after 2nd resize
────────────────
Memory usage: 311.1 MB
Object growth:


after 1 additional resize(s)
────────────────────────────
Memory usage: 316.8 MB
Object growth:


after 2 additional resize(s)
────────────────────────────
Memory usage: 295.4 MB
Object growth:


after 3 additional resize(s)
────────────────────────────
Memory usage: 295.4 MB
Object growth:


after 4 additional resize(s)
────────────────────────────
Memory usage: 326.2 MB
Object growth:


after 5 additional resize(s)
────────────────────────────
Memory usage: 327.3 MB
Object growth:


after 6 additional resize(s)
────────────────────────────
Memory usage: 1.1 GB
Object growth:


after 7 additional resize(s)
────────────────────────────
Memory usage: 324.1 MB
Object growth:


after 8 additional resize(s)
────────────────────────────
Memory usage: 321.6 MB
Object growth:


after 9 additional resize(s)
────────────────────────────
Memory usage: 295.4 MB
Object growth:


after 10 additional resize(s)
─────────────────────────────
Memory usage: 295.4 MB
Object growth:


Press enter to continue 
Graph written to /tmp/objgraph-cy0orb9r.dot (3 nodes)
Image generated as <_io.BytesIO object at 0x7fbc78071780> chain.png
Graph written to /tmp/objgraph-hyputy63.dot (3 nodes)
Image generated as <_io.BytesIO object at 0x7fbc7800ed58> chain.png