I am using a Python class that needs to keep track of callback functions/methods, and the instance of this class needs to have a reference (amongst other places) in one of the classes that provides a callback method. However, this is leading to a memory leak due to circular referencing of objects.
This is easier to explain in code:
import psutil
import numpy as np
process = psutil.Process()
def show_mem(title):
mem = process.memory_info()
print(title + ':')
print('Real: {0:.1f}Mb'.format(mem.rss / 1024**2))
print('Virtual: {0:.1f}Mb'.format(mem.vms / 1024**2))
print()
class CallbackTracker(object):
def __init__(self):
self.callbacks = []
def add_callback(self, callback):
self.callbacks.append(callback)
class TestClass(object):
def __init__(self):
self.array = np.random.random(100000000)
self.ct = CallbackTracker()
self.ct.add_callback(self.callback)
def callback(self):
pass
def run():
t = TestClass()
show_mem('Before run()')
run()
show_mem('After run()')
When run, this returns:
Before run():
Real: 22.4Mb
Virtual: 2451.9Mb
After run():
Real: 785.4Mb
Virtual: 3214.8Mb
The array is there to just make the memory usage more obvious. If I remove the following line:
self.ct.add_callback(self.callback)
Then the memory leak goes away:
Before run():
Real: 22.3Mb
Virtual: 2451.9Mb
After run():
Real: 22.3Mb
Virtual: 2451.9Mb
What is the best way to avoid this? I have tried using weak references, but the issue is that the callback function appears to be dereferenced immediately after the class is instantiated:
import numpy as np
from weakref import proxy
class CallbackTracker(object):
def __init__(self):
self.callbacks = []
def remove_callback(self, pcallback):
print('removing callback')
self.callbacks.remove(pcallback)
def add_callback(self, callback):
self.callbacks.append(proxy(callback, self.remove_callback))
class TestClass(object):
def __init__(self):
self.array = np.random.random(100000000)
self.ct = CallbackTracker()
self.ct.add_callback(self.callback)
def callback(self):
pass
t = TestClass()
print('done with TestClass()')
gives:
removing callback
done with TestClass()
which is no good because the callbacks (in the real code) would then not work. It's not clear to me why self.callback
gets dereferenced before the end of the script. What is the best solution to my original problem, and is there a way to fix the weakref solution?